#include "socks4a.h"

#include <memory>

#include "muduo/base/Logging.h"
#include "muduo/net/Endian.h"
using namespace muduo;
using namespace muduo::net;

Socks4a::Socks4a(EventLoop *loop, const InetAddress &listenAddr)
    : TcpServer(loop, listenAddr, "Socks4"), resolver_(loop)
{
  this->setConnectionCallback(std::bind(&Socks4a::onServerConnection, this, _1));
  this->setMessageCallback(std::bind(&Socks4a::onServerMessage, this, _1, _2, _3));
}
/**
 * @brief 处理收到客户端发来连接的情况
 *
 * @param conn
 */
void Socks4a::onServerConnection(const TcpConnectionPtr &conn)
{
  LOG_DEBUG << conn->name() << (conn->connected() ? " UP" : " DOWN");
  if (conn->connected()) {
    // 禁用Nagle算法，避免连续发包出现延迟，这对编写低延迟网络服务很重要
    conn->setTcpNoDelay(true);
  } else {
    auto it = tunnels_.find(conn->name());
    if (it != tunnels_.end()) {
      it->second->disconnect();
      tunnels_.erase(it);
    }
  }
}

/**
 * @brief 处理客户端来的消息，以及第一个消息保存conn
 *
 * @param conn 与客户端的连接
 * @param buf 客户端发来的buf
 */
void Socks4a::onServerMessage(const TcpConnectionPtr &conn, Buffer *buf, muduo::Timestamp)
{
  LOG_DEBUG << conn->name() << " " << buf->readableBytes();
  if (tunnels_.find(conn->name()) == tunnels_.end()) {
    size_t len = buf->readableBytes();
    if (len > 128) {
      conn->shutdown();
    } else if (len > 8) {
      codecSocks4(conn, buf, buf->peek(), len);
    }
  } else if (!conn->getContext().empty()) {
    const TcpConnectionPtr &clientConn =
        boost::any_cast<const TcpConnectionPtr &>(conn->getContext());
    clientConn->send(buf);
  }
}

/**
 * @brief 解析socks4协议，目前只支持41开头，即socks4*
 * connect，成功就返回地址对应的unique_ptr，否则返回nullptr
 *
 * @param data 需要解析的字符串
 * @param len  字符串的长度
 * @param where  是返回值，最后会指向域名的后一位，即socks4 中的 NULL，失败则指向nullptr
 * @return void
 */
void Socks4a::codecSocks4(const TcpConnectionPtr &conn, Buffer *buf, const char *data, size_t len)
{
  const char *begin = data + 8;
  const char *end = data + len;
  const char *where = std::find(begin, end, '\0');
  if (where != end) {
    char ver = data[0];
    char cmd = data[1];
    if (ver != 4 || cmd != 1) {
      conn->shutdown();
      return;
    }
    const void *port = data + 2;
    const void *ip = data + 4;
    sockaddr_in addr;
    memZero(&addr, sizeof addr);
    addr.sin_family = AF_INET;
    addr.sin_port = *static_cast<const in_port_t *>(port);
    addr.sin_addr.s_addr = *static_cast<const uint32_t *>(ip);
    // 前24位都是0，那么就< 256,也就是socks4a，否则就是socks4
    bool socks4a = sockets::networkToHost32(addr.sin_addr.s_addr) < 256;

    if (socks4a) {
      const char *endOfHostName = std::find(where + 1, end, '\0');

      if (endOfHostName != end) {
        std::string hostname = where + 1;
        where = endOfHostName;
        LOG_INFO << "Socks4a host name -------- " << hostname;
        InetAddress tmp;
        // 获取域名对应的ip
        resolver_.resolve(hostname, [conn, buf, where, addr,
                                     this](const InetAddress &ip) mutable {
          // 是socks4a
          if (ip.ipv4NetEndian() != 0) {
            addr.sin_addr.s_addr = ip.ipv4NetEndian();
            farwardConnect(conn, buf, where, muduo::net::InetAddress(addr));
          } else
            conn->shutdown();
        });
        // 获取域名对应的ip
      } else
        conn->shutdown();  // 说明不是socks4*
    } else
      farwardConnect(conn, buf, where, InetAddress(addr));  // 这里要调用直接ip的 socks4
  } else
    conn->shutdown();  // 说明不是socks4*
}

void Socks4a::farwardConnect(const TcpConnectionPtr &conn, Buffer *buf, const char *where,
                             const InetAddress &serverAddr)
{
  // socks4 or 4a
  LOG_INFO << "Socks4a host ip --------" << serverAddr.toIp();
  // 开启通道
  auto tunnel = std::make_shared<Tunnel>(this->getLoop(), serverAddr, conn);
  tunnel->setup();
  tunnel->connect();
  tunnels_[conn->name()] = tunnel;
  buf->retrieveUntil(where + 1);
  // 发送回应
  char response[] = "\000\x5aUVWXYZ";
  auto sin_port{serverAddr.portNetEndian()};
  auto s_addr{serverAddr.ipv4NetEndian()};
  memcpy(response + 2, &sin_port, 2);
  memcpy(response + 4, &s_addr, 4);
  conn->send(response, 8);
}


void Socks4a::start()
{
  TcpServer::start();
  resolver_.start();
}
